# -*- coding: utf-8 -*-
"""
A股股票分钟级数据同步模块
从akshare获取股票分钟级交易数据，并同步到数据库
"""
import logging
import pandas as pd
import akshare as ak
from datetime import datetime, timedelta
from typing import List, Dict, Any, Optional, Tuple
from database import DatabaseManager, StockMinuteDAO
from config import DATABASE_CONFIG, SYNC_CONFIG

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class StockMinuteSync:
    """A股股票分钟级数据同步器"""
    
    def __init__(self):
        self.db_manager = DatabaseManager(DATABASE_CONFIG)
        self.minute_dao = StockMinuteDAO(self.db_manager)
        self.batch_size = SYNC_CONFIG.get('batch_size', 1000)
        self.request_delay = SYNC_CONFIG.get('request_delay', 2.0)  # 分钟级数据请求间隔
        self.max_retries = 3  # 最大重试次数
    
    def get_stock_codes_with_market(self) -> List[Tuple[str, str]]:
        """获取所有活跃股票代码及其市场信息"""
        try:
            sql = """
            SELECT stock_code, exchange 
            FROM new_stock_info 
            WHERE status = 'active' 
            ORDER BY stock_code
            """
            results = self.db_manager.execute_query(sql)
            stock_market_list = []
            for row in results:
                stock_code = row['stock_code']
                exchange = row['exchange']
                # 将交易所代码转换为市场代码
                if exchange == 'SZSE':
                    market = 'sz'
                elif exchange == 'SHSE':
                    market = 'sh'
                else:
                    continue  # 跳过未知交易所
                stock_market_list.append((stock_code, market))
            
            logger.info(f"获取到 {len(stock_market_list)} 只活跃股票")
            return stock_market_list
        except Exception as e:
            logger.error(f"获取股票代码失败: {e}")
            raise
    
    def get_latest_trade_time(self, stock_code: str, market: str) -> Optional[datetime]:
        """获取指定股票和市场的最新交易时间"""
        return self.minute_dao.get_latest_trade_time(stock_code, market)
    
    def get_minute_data(self, stock_code: str, market: str, 
                       start_time: datetime = None, end_time: datetime = None) -> pd.DataFrame:
        """获取单只股票的分钟级数据"""
        try:
            # 如果没有指定时间，默认获取最近1小时的数据
            if not start_time:
                start_time = datetime.now() - timedelta(hours=1)
            if not end_time:
                end_time = datetime.now()
            
            # akshare的分钟级数据接口需要使用不带后缀的股票代码
            symbol = stock_code
            
            logger.debug(f"正在获取 {symbol} 的分钟级数据，时间范围: {start_time} - {end_time}")
            
            # 使用akshare获取分钟级数据，带重试机制
            df = None
            for attempt in range(self.max_retries):
                try:
                    # 使用akshare的分钟级数据接口
                    df = ak.stock_zh_a_hist_min_em(
                        symbol=symbol,
                        period="1",  # 1分钟
                        start_date=start_time.strftime('%Y%m%d %H:%M:%S'),
                        end_date=end_time.strftime('%Y%m%d %H:%M:%S'),
                        adjust="qfq"  # 前复权
                    )
                    break  # 成功获取数据，跳出重试循环
                except Exception as e:
                    error_msg = str(e)
                    # 检查是否是连接被拒绝的错误
                    if "Connection aborted" in error_msg or "Remote end closed connection" in error_msg:
                        logger.error(f"检测到连接被拒绝错误，立即终止同步: {e}")
                        raise SystemExit("连接被拒绝，程序终止")
                    
                    if attempt < self.max_retries - 1:
                        logger.warning(f"获取 {symbol} 分钟级数据失败，第 {attempt + 1} 次重试: {e}")
                        import time
                        time.sleep(3)  # 重试前等待3秒
                    else:
                        raise e  # 最后一次重试失败，抛出异常
            
            if df.empty:
                logger.warning(f"股票 {symbol} 没有获取到分钟级数据")
                return pd.DataFrame()
            
            # 重命名列
            df.rename(columns={
                "时间": "trade_time",
                "开盘": "open",
                "收盘": "close",
                "最高": "high",
                "最低": "low",
                "成交量": "volume",
                "成交额": "amount"
            }, inplace=True)
            
            # 添加股票代码和市场信息
            df['stock_code'] = stock_code
            df['market'] = market
            
            # 选择需要的列
            required_columns = ['stock_code', 'market', 'trade_time', 'open', 'close', 'high', 'low', 'volume', 'amount']
            df = df[required_columns]
            
            # 处理数据类型
            df['trade_time'] = pd.to_datetime(df['trade_time'])
            df['open'] = pd.to_numeric(df['open'], errors='coerce')
            df['close'] = pd.to_numeric(df['close'], errors='coerce')
            df['high'] = pd.to_numeric(df['high'], errors='coerce')
            df['low'] = pd.to_numeric(df['low'], errors='coerce')
            df['volume'] = pd.to_numeric(df['volume'], errors='coerce')
            df['amount'] = pd.to_numeric(df['amount'], errors='coerce')
            
            # 删除无效数据
            df = df.dropna(subset=['open', 'close', 'high', 'low', 'volume', 'amount'])
            
            logger.debug(f"成功获取 {symbol} 的分钟级数据: {len(df)} 条记录")
            return df
            
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 分钟级数据失败: {e}")
            return pd.DataFrame()
    
    def sync_single_stock_minute(self, stock_code: str, market: str, 
                                max_hours: int = 24) -> Dict[str, Any]:
        """同步单只股票的分钟级数据"""
        result = {
            'stock_code': stock_code,
            'market': market,
            'success': False,
            'records_count': 0,
            'error': None
        }
        
        try:
            # 获取该股票的最新交易时间
            latest_time = self.get_latest_trade_time(stock_code, market)
            
            # 确定同步的时间范围
            end_time = datetime.now()
            
            if latest_time:
                # 从最新时间的下1分钟开始同步
                start_time = latest_time + timedelta(minutes=1)
                
                # 如果最新时间就是现在，跳过
                if start_time >= end_time:
                    result['success'] = True
                    result['records_count'] = 0
                    logger.info(f"股票 {stock_code} 分钟级数据已是最新，跳过同步")
                    return result
            else:
                # 如果没有历史数据，从指定小时前开始同步
                start_time = datetime.now() - timedelta(hours=max_hours)
            
            logger.info(f"同步股票 {stock_code} 分钟级数据，时间范围: {start_time} - {end_time}")
            
            # 获取分钟级数据
            df = self.get_minute_data(stock_code, market, start_time, end_time)
            
            if df.empty:
                result['success'] = True
                result['records_count'] = 0
                logger.warning(f"股票 {stock_code} 没有获取到新的分钟级数据")
                return result
            
            # 立即保存到数据库
            batch_data = df.to_dict('records')
            synced_count = self.minute_dao.batch_insert_or_update_minute_data(batch_data)
            
            result['success'] = True
            result['records_count'] = synced_count
            logger.info(f"股票 {stock_code} 分钟级数据同步完成，新增 {synced_count} 条记录")
            
        except Exception as e:
            result['error'] = str(e)
            logger.error(f"股票 {stock_code} 分钟级数据同步失败: {e}")
        
        return result
    
    def sync_minute_data(self, stock_market_list: List[Tuple[str, str]] = None, 
                        max_hours: int = 24) -> Dict[str, Any]:
        """同步股票分钟级数据（增量同步）"""
        start_time = datetime.now()
        result = {
            'success': False,
            'total_stocks': 0,
            'successful_stocks': 0,
            'failed_stocks': 0,
            'total_records': 0,
            'synced_records': 0,
            'error': None,
            'duration': 0,
            'failed_list': []
        }
        
        try:
            # 获取股票代码列表
            if not stock_market_list:
                stock_market_list = self.get_stock_codes_with_market()
            
            result['total_stocks'] = len(stock_market_list)
            logger.info(f"开始增量同步 {len(stock_market_list)} 只股票的分钟级数据...")
            
            successful_count = 0
            failed_count = 0
            failed_list = []
            total_synced_records = 0
            
            # 逐个同步股票数据
            for i, (stock_code, market) in enumerate(stock_market_list, 1):
                try:
                    logger.info(f"正在处理第 {i}/{len(stock_market_list)} 只股票: {stock_code}.{market.upper()}")
                    
                    # 同步单只股票
                    stock_result = self.sync_single_stock_minute(stock_code, market, max_hours)
                    
                    if stock_result['success']:
                        successful_count += 1
                        total_synced_records += stock_result['records_count']
                        if stock_result['records_count'] > 0:
                            logger.info(f"✓ {stock_code}.{market.upper()}: 新增 {stock_result['records_count']} 条记录")
                        else:
                            logger.info(f"✓ {stock_code}.{market.upper()}: 数据已是最新")
                    else:
                        failed_count += 1
                        failed_list.append(f"{stock_code}.{market.upper()}")
                        logger.error(f"✗ {stock_code}.{market.upper()}: 同步失败 - {stock_result['error']}")
                    
                    # 添加请求延迟，避免频率限制
                    if self.request_delay > 0:
                        import time
                        time.sleep(self.request_delay)
                        
                except Exception as e:
                    failed_count += 1
                    failed_list.append(f"{stock_code}.{market.upper()}")
                    logger.error(f"✗ {stock_code}.{market.upper()}: 处理失败 - {e}")
                    continue
            
            result['success'] = True
            result['successful_stocks'] = successful_count
            result['failed_stocks'] = failed_count
            result['total_records'] = total_synced_records
            result['synced_records'] = total_synced_records
            result['failed_list'] = failed_list
            
            logger.info(f"分钟级数据同步完成！成功: {successful_count} 只，失败: {failed_count} 只，共同步 {total_synced_records} 条记录")
            
        except Exception as e:
            result['error'] = str(e)
            logger.error(f"分钟级数据同步失败: {e}")
        
        finally:
            end_time = datetime.now()
            result['duration'] = (end_time - start_time).total_seconds()
            logger.info(f"同步耗时: {result['duration']:.2f} 秒")
        
        return result
    
    def sync_specific_stocks_minute(self, stock_codes: List[str], 
                                  max_hours: int = 24) -> Dict[str, Any]:
        """同步指定股票的分钟级数据"""
        stock_market_list = []
        
        for stock_code in stock_codes:
            # 根据股票代码判断市场
            if stock_code.startswith('000') or stock_code.startswith('002') or stock_code.startswith('300'):
                market = 'sz'
            elif stock_code.startswith('600') or stock_code.startswith('601') or stock_code.startswith('603') or stock_code.startswith('605') or stock_code.startswith('688'):
                market = 'sh'
            else:
                logger.warning(f"无法识别股票代码 {stock_code} 的市场，跳过")
                continue
            
            stock_market_list.append((stock_code, market))
        
        return self.sync_minute_data(stock_market_list, max_hours)
    
    def get_sync_summary(self) -> Dict[str, Any]:
        """获取同步摘要信息"""
        return self.minute_dao.get_data_statistics()

def main():
    """主函数"""
    try:
        # 创建同步器
        sync = StockMinuteSync()
        
        # 执行同步
        result = sync.sync_minute_data()
        
        # 打印结果
        if result['success']:
            print("=" * 60)
            print("A股分钟级数据同步成功！")
            print(f"处理股票数: {result['total_stocks']} 只")
            print(f"成功股票数: {result['successful_stocks']} 只")
            print(f"失败股票数: {result['failed_stocks']} 只")
            print(f"总记录数: {result['total_records']} 条")
            print(f"同步记录数: {result['synced_records']} 条")
            print(f"同步耗时: {result['duration']:.2f} 秒")
            if result['failed_list']:
                print(f"失败股票: {', '.join(result['failed_list'][:10])}{'...' if len(result['failed_list']) > 10 else ''}")
            print("=" * 60)
        else:
            print("=" * 60)
            print("A股分钟级数据同步失败！")
            print(f"错误信息: {result['error']}")
            print("=" * 60)
        
        # 显示同步摘要
        summary = sync.get_sync_summary()
        if 'error' not in summary:
            print("\n数据库当前状态:")
            print(f"总记录数: {summary['total_records']}")
            print(f"股票数: {summary['total_stocks']}")
            print(f"最新时间: {summary['latest_time']}")
            print(f"最早时间: {summary['earliest_time']}")
        
    except Exception as e:
        logger.error(f"程序执行失败: {e}")
        print(f"程序执行失败: {e}")

if __name__ == "__main__":
    main()
